import argparse
import os

import numpy
import torch
from torch.utils.data import DataLoader
from ConvNetQuake_pytorch.constant import EVENT_CLUSTER, TRAIN_PATH
from models import ConvNetQuake
from eq_dataset import EarthquakeDatasetCache
import numpy as np
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument("--n_cluster", default=1)
args = parser.parse_args()

device = "cuda" if torch.cuda.is_available() else "cpu"
model = ConvNetQuake(args).to(device)
model_path = r'..\data\mytrain\backup\20220404_1\e12_c0.9542413806752958.pth'
model.load_state_dict(torch.load(model_path))
model.eval()

test_path = os.path.join(TRAIN_PATH, "test.csv")
test_dataset = EarthquakeDatasetCache(test_path, True)
test_dataset.init_cache(2)
print(len(test_dataset))
test_dataset_loader = DataLoader(dataset=test_dataset, batch_size=10, shuffle=False)

predictions = torch.tensor([])
labels = torch.tensor([])

for data, label in tqdm(test_dataset_loader):
    data, label = data.to(device), label.to(device)
    prediction = model(data)
    predictions = torch.cat((predictions, torch.argmax(prediction, dim=1).cpu()))
    labels = torch.cat((labels, label.cpu()))
    # print("label\t = {}".format(label))
    # print("pred\t = {}".format(torch.argmax(prediction, dim=1)))
    # print()
predictions = predictions.numpy().astype(np.int32)
labels = labels.numpy().astype(np.int32)

confusion_matrix = numpy.zeros((2, 2))
for t, p in zip(labels, predictions):
    confusion_matrix[t, p] += 1
print(confusion_matrix)
tp = confusion_matrix[1, 1]
fp = confusion_matrix[0, 1]
fn = confusion_matrix[1, 0]
tn = confusion_matrix[0, 0]
acc = (tn + tp) / confusion_matrix.sum()
recall = tp / (fn + tp)
precision = tp / (tp + fp)
print("acc = {}, recall = {}, precision = {}".format(acc, recall, precision))
